import os
import shutil
import sys
from glob import glob

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
# 构造结构化数据集
class CompetiontionDataset(Dataset):
    def __init__(self):
        self.image_names_path = glob(os.path.join('Img', '*.jpg'))
        self.to_tensor = transforms.ToTensor()


    def __len__(self):
        return len(self.image_names_path)

    def __getitem__(self, idx):
        image_name = self.image_names_path[idx][-10:-4]
        image = Image.open(os.path.join('Img', f'{image_name}.jpg'))
        return self.to_tensor(image), image_name

def preprocess_path():
    """
    处理目录路径。
    此函数旨在确保目录'./c_test_predictions'在程序运行前被创建。
    如果目录已存在，则先尝试删除再重新创建。
    如果目录不存在，则直接创建。
    """
    check_status = True
    if os.path.exists('../c_test_predictions'):
        try:
            shutil.rmtree('../c_test_predictions')
            print('存在c_test_predictions目录，已经删除')
        except:
            print('【Error】删除c_test_predictions目录失败')
            check_status = False
    try:
        os.mkdir('../c_test_predictions')
        print('创建c_test_predictions目录成功')
    except:
        print('【Error】创建c_test_predictions目录失败')
        check_status = False

    if not os.path.exists('../model.py'):
        print('【Error】未找到model.py文件，请将model.py文件放入c_test_predictions同级目录')
        check_status = False

    if not os.path.exists('../model.pth'):
        print('【Error】未找到model.pth文件，请将model.pth文件放入c_test_predictions同级目录')
        check_status = False

    if not os.path.exists('../关键指标数据文档.txt'):
        print('【Error】未找到关键指标数据文档.txt文件，请将关键指标数据文档.txt文件放入c_test_predictions同级目录')
        check_status = False

    if not os.path.exists('../算法描述文档.pdf'):
        print('【Error】未找到算法描述文档.pdf文件，请将算法描述文档.pdf文件放入c_test_predictions同级目录')
        check_status = False

    return check_status

def convert_pred_to_mask(pred):
    """
    将预测张量转换为预测掩膜。

    参数:
    - pred (torch.Tensor): 形状为 (batch_size, num_classes, h, w) 的预测张量。

    返回:
    - torch.Tensor: 形状为 (batch_size, h, w) 的预测掩膜。
    """
    # 选择每个像素的最大概率类别
    _, pred_mask = torch.max(pred, dim=1)
    pred_mask = pred_mask.to(torch.long)
    return pred_mask

def eval_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = self_net().to(device)
    model.load_state_dict(torch.load('../model.pth'))
    model.eval()
    test_dataset = CompetiontionDataset()
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
    for i, (input, image_name) in enumerate(test_loader):
        input = input.to(device)
        pred = model(input)
        pred_mask = convert_pred_to_mask(pred)
        pred_musk_np = pred_mask.cpu().numpy()
        np.save(f'../c_test_predictions/c_prediction_{image_name[0]}.npy', pred_musk_np)


if __name__ == '__main__':
    sys.path.append("..") #目的是为了导入 model.py 中的模型 所以把 上级目录加入  self_net可能提示 未解析引用，其实是引用了。
    from model import *
    if preprocess_path():  #测试前做环境预处理
        eval_model()  # 生成npy文件
        print('验证完成，你现在可以提交赛克了')
    else:
        print('环境预处理失败，请检查缺失文件')








